
#include <cassert>
#include <cstdlib>

#include <DXGI.h>
#include <D3D10_1.h>
#include <D3DX11.h>
#include <D3DX10math.h>
#include <D3DX11async.h>

#include "Core\RenderTarget.h"
#include "Math\Matrix4x4.h"
#include "Scene\Light.h"
#include "Shaders\DeferredLightingShader.h"
#include "Tools\RGBColor.h"

#pragma comment(lib, "D3D11.lib")
#pragma comment(lib, "D3DX11.lib")


ID3D10Blob* DeferredLightingShader::s_vshader_buffer = NULL;
ID3D10Blob* DeferredLightingShader::s_pshader_buffer = NULL;


BOOL DeferredLightingShader::s_loaded = FALSE;


DeferredLightingShader::DeferredLightingShader()
	: m_vshader(NULL),
	  m_pshader(NULL),
	  m_layout(NULL),
	  m_binded(FALSE),
	  m_light_buffer(NULL)
{
}


DeferredLightingShader::~DeferredLightingShader()
{
	// Unbind if needed.
	if(m_binded)
		Unbind();
}


BOOL DeferredLightingShader::IsLoaded()
{
	return s_loaded;
}


BOOL DeferredLightingShader::Load()
{
	assert(!s_vshader_buffer);
	assert(!s_pshader_buffer);
	assert(!s_loaded);

	ID3D10Blob* error_buffer = NULL;
	LPCSTR compile_error = NULL;
	HRESULT result;

	// Load and compile vertex shader.
	result = D3DX11CompileFromFile(TEXT("Shaders\\RT2Screen.vs"),
								   NULL, NULL, "Main", "vs_5_0",
								   D3D10_SHADER_ENABLE_STRICTNESS, 0, NULL,
								   &s_vshader_buffer, &error_buffer, NULL);

	if(FAILED(result))
		compile_error = static_cast< LPCSTR >(error_buffer->GetBufferPointer());

	assert(SUCCEEDED(result));

	// Load and compile pixel shader.
	result = D3DX11CompileFromFile(TEXT("Shaders\\DeferredLighting.ps"),
								NULL, NULL, "Main", "ps_5_0",
								D3D10_SHADER_ENABLE_STRICTNESS, 0, NULL,
								&s_pshader_buffer, &error_buffer, NULL);

	if(FAILED(result))
		compile_error = static_cast< LPCSTR >(error_buffer->GetBufferPointer());

	assert(SUCCEEDED(result));

	// Change state to loaded.
	s_loaded = TRUE;

	return TRUE;
}


VOID DeferredLightingShader::Unload()
{
	assert(s_vshader_buffer);
	assert(s_pshader_buffer);
	assert(s_loaded);

	s_vshader_buffer->Release();
	s_vshader_buffer = NULL;

	s_pshader_buffer->Release();
	s_pshader_buffer = NULL;

	// Change state to unloaded.
	s_loaded = FALSE;
	s_vshader_buffer = NULL;
	s_pshader_buffer = NULL;
}


BOOL DeferredLightingShader::IsBinded() CONST
{
	return m_binded;
}


BOOL DeferredLightingShader::Bind()
{
	assert(s_vshader_buffer);
	assert(s_pshader_buffer);
	assert(!m_layout);
	assert(!m_vshader);
	assert(!m_pshader);
	assert(s_loaded);
	assert(!m_binded);

	HRESULT result;

	// Create the vertex shader buffer.
	result = GetDevice()->CreateVertexShader(s_vshader_buffer->GetBufferPointer(),
										     s_vshader_buffer->GetBufferSize(),
										     NULL,
									         &m_vshader);
	
	assert(SUCCEEDED(result));

	// Create the pixel shader buffer.
	result = GetDevice()->CreatePixelShader(s_pshader_buffer->GetBufferPointer(),
										    s_pshader_buffer->GetBufferSize(),
											NULL,
											&m_pshader);

	// Setup the input layout description.
	const UINT input_count = 2;
	D3D11_INPUT_ELEMENT_DESC input_desc[input_count];

	input_desc[0].SemanticName = "POSITION";
	input_desc[0].SemanticIndex = 0;
	input_desc[0].Format = DXGI_FORMAT_R32G32B32A32_FLOAT;
	input_desc[0].InputSlot = 0;
	input_desc[0].AlignedByteOffset = 0;
	input_desc[0].InputSlotClass = D3D11_INPUT_PER_VERTEX_DATA;
	input_desc[0].InstanceDataStepRate = 0;

	input_desc[1].SemanticName = "TEXCOORD";
	input_desc[1].SemanticIndex = 0;
	input_desc[1].Format = DXGI_FORMAT_R32G32_FLOAT;
	input_desc[1].InputSlot = 0;
	input_desc[1].AlignedByteOffset = D3D11_APPEND_ALIGNED_ELEMENT;
	input_desc[1].InputSlotClass = D3D11_INPUT_PER_VERTEX_DATA;
	input_desc[1].InstanceDataStepRate = 0;

	// Create the input layout.
	result = GetDevice()->CreateInputLayout(input_desc,
										    input_count,
									        s_vshader_buffer->GetBufferPointer(),
										    s_vshader_buffer->GetBufferSize(),
									        &m_layout);

	assert(SUCCEEDED(result));

	// Setup dynamic buffers.
	D3D11_BUFFER_DESC dynamic_desc;

	dynamic_desc.Usage = D3D11_USAGE_DYNAMIC;
	dynamic_desc.BindFlags = D3D11_BIND_CONSTANT_BUFFER;
	dynamic_desc.CPUAccessFlags = D3D11_CPU_ACCESS_WRITE;
	dynamic_desc.MiscFlags = 0;
	dynamic_desc.StructureByteStride = 0;

	dynamic_desc.ByteWidth = sizeof(LightBuffer);

	result = GetDevice()->CreateBuffer(&dynamic_desc, NULL, &m_light_buffer);
	assert(SUCCEEDED(result));

	// Change state to binded.
	m_binded = TRUE;

	return m_binded;
}


VOID DeferredLightingShader::Unbind()
{
	assert(m_layout);
	assert(m_vshader);
	assert(m_pshader);
	assert(m_binded);

	m_vshader->Release();
	m_vshader = NULL;

	m_pshader->Release();
	m_pshader = NULL;

	m_layout->Release();
	m_layout = NULL;

	m_light_buffer->Release();
	m_light_buffer = NULL;
	
	// Change state to unbinded.
	m_binded = FALSE;
}


VOID DeferredLightingShader::Enable()
{
	assert(m_layout);
	assert(m_vshader);
	assert(m_pshader);
	assert(m_binded);

	// Enabling.
	GetContext()->IASetInputLayout(m_layout);
	GetContext()->VSSetShader(m_vshader, NULL, 0);
	GetContext()->PSSetShader(m_pshader, NULL, 0);
}


BOOL DeferredLightingShader::SetLightBuffer(CONST Light& light, CONST Matrix4x4& world)
{
	assert(m_layout);
	assert(m_vshader);
	assert(m_pshader);
	assert(m_binded);

	// Map lights.
	D3D11_MAPPED_SUBRESOURCE mapped;

	HRESULT result = GetContext()->Map(m_light_buffer,
								       0,
								       D3D11_MAP_WRITE_DISCARD,
								       0,
								       &mapped);

	assert(SUCCEEDED(result));

	LightBuffer* ptr = reinterpret_cast< LightBuffer* >(mapped.pData);

	Vector4 position = world * light.GetPosition();
	ptr->position[0] = position.GetX();
	ptr->position[1] = position.GetY();
	ptr->position[2] = position.GetZ();
	ptr->position[3] = position.GetW();

	Vector4 direction = world * light.GetDirection();
	ptr->direction[0] = direction.GetX();
	ptr->direction[1] = direction.GetY();
	ptr->direction[2] = direction.GetZ();

	ptr->angle = light.GetAngle();

	ptr->radius = light.GetRadius();

	RGBColor ambient = light.GetAmbient();
	ptr->ambient[0] = ambient.GetR();
	ptr->ambient[1] = ambient.GetG();
	ptr->ambient[2] = ambient.GetB();

	RGBColor diffuse = light.GetDiffuse();
	ptr->diffuse[0] = diffuse.GetR();
	ptr->diffuse[1] = diffuse.GetG();
	ptr->diffuse[2] = diffuse.GetB();

	RGBColor specular = light.GetSpecular();
	ptr->specular[0] = specular.GetR();
	ptr->specular[1] = specular.GetG();
	ptr->specular[2] = specular.GetB();

	ptr->constant_attenuation = light.GetConstantAttenuation();

	ptr->linear_attenuation = light.GetLinearAttenuation();

	ptr->quadratic_attenuation = light.GetQuadraticAttenuation();
		
	ptr->exponent = light.GetExponent();

	GetContext()->Unmap(m_light_buffer, 0);

	// Set lights constant buffer.
	GetContext()->PSSetConstantBuffers(0, 1, &m_light_buffer);

	return TRUE;
}


BOOL DeferredLightingShader::SetGBuffer(RenderTarget* CONST gbuffer[GBUFFER_RT_COUNT])
{
	assert(m_vshader);
	assert(m_pshader);
	assert(m_binded);



	// Unset the source.
	if(!m_gbuffer)
	{
		for(UINT8 i = 0; i < GBUFFER_RT_COUNT; ++i)
		{
			m_gbuffer[i] = NULL;
			GetContext()->PSSetShaderResources(i, 1, NULL);
		}

		return TRUE;
	}

	for(UINT8 i = 0; i < GBUFFER_RT_COUNT; ++i)
	{
		m_gbuffer[i] = gbuffer[i];
		
		ID3D11ShaderResourceView* view = m_gbuffer[i]->GetResource();
		GetContext()->PSSetShaderResources(i, 1, &view);
	}

	return TRUE;
}